"""
@Description :   评价指标计算
@Author      :   tqychy 
@Time        :   2025/08/23 17:10:45
"""

import numpy as np


def iou(bbox1, bbox2):
    """
    计算两个边界框的交并比 (IoU)。

    参数:
        bbox1 (list or tuple): 第一个边界框，格式为 [x_min, y_min, width, height]。
        bbox2 (list or tuple): 第二个边界框，格式为 [x_min, y_min, width, height]。

    返回:
        float: IoU 值，范围在 [0, 1] 之间。如果两个框没有交集，则返回 0.0。
    """
    x_min1, y_min1, w1, h1 = bbox1
    x_min2, y_min2, w2, h2 = bbox2

    # 计算交集区域的坐标
    # 左上角 (x, y) 是两个框左上角坐标的最大值
    inter_x_min = max(x_min1, x_min2)
    inter_y_min = max(y_min1, y_min2)

    # 右下角 (x, y) 是两个框右下角坐标的最小值
    # 右下角坐标 = 左上角坐标 + 宽度/高度
    inter_x_max = min(x_min1 + w1, x_min2 + w2)
    inter_y_max = min(y_min1 + h1, y_min2 + h2)

    # 计算交集的宽度和高度
    # 如果计算出的宽度或高度为负数，说明没有交集，应设为 0
    inter_width = max(0, inter_x_max - inter_x_min)
    inter_height = max(0, inter_y_max - inter_y_min)

    # 计算交集面积
    inter_area = inter_width * inter_height

    # 计算两个框各自的面积
    area1 = w1 * h1
    area2 = w2 * h2

    # 计算并集面积
    union_area = area1 + area2 - inter_area

    # 防止除以零（当两个框面积都为0时，union_area为0）
    if union_area == 0:
        return 0.0

    # 计算并返回 IoU
    iou = inter_area / union_area
    return iou

def precision_recall(predictions, ground_truths, classes, iou_threshold):
    """
    计算给定IoU阈值下每个类别的精确度和召回率。
    返回:
        精确度和召回率的字典元组
    """
    precision = {c: [] for c in classes}
    recall = {c: [] for c in classes}
    
    for pred, gt in zip(predictions, ground_truths):
        pred_bboxes = pred['bboxes']
        pred_cats = pred['categories']
        pred_scores = pred['scores']
        gt_bboxes = gt['bboxes']
        gt_cats = gt['categories']
        
        # 按分数排序预测
        sorted_idx = np.argsort(pred_scores)[::-1]
        pred_bboxes = pred_bboxes[sorted_idx]
        pred_cats = pred_cats[sorted_idx]
        
        for c in classes:
            gt_mask = gt_cats == c
            pred_mask = pred_cats == c
            
            gt_class_bboxes = gt_bboxes[gt_mask]
            pred_class_bboxes = pred_bboxes[pred_mask]
            
            if len(gt_class_bboxes) == 0 and len(pred_class_bboxes) == 0:
                continue
                
            matches = []
            matched_gt = set()
            for i, pred_box in enumerate(pred_class_bboxes):
                max_iou = 0
                best_gt_idx = -1
                
                for j, gt_box in enumerate(gt_class_bboxes):
                    if j in matched_gt:
                        continue
                    iou_score = iou(pred_box, gt_box)
                    if iou_score > max_iou:
                        max_iou = iou_score
                        best_gt_idx = j
                        
                if max_iou >= iou_threshold and best_gt_idx != -1:
                    matches.append(1)
                    matched_gt.add(best_gt_idx)
                else:
                    matches.append(0)
                    
            tp = sum(matches)
            fp = len(pred_class_bboxes) - tp
            fn = len(gt_class_bboxes) - tp
            
            prec = tp / (tp + fp) if (tp + fp) > 0 else 0
            rec = tp / (tp + fn) if (tp + fn) > 0 else 0
            
            precision[c].append(prec)
            recall[c].append(rec)
            
    return precision, recall


def map(predictions, ground_truths, classes, iou_thresholds):
    """
    计算所有类别和IoU阈值的平均精度均值（mAP）。

    返回:
        包含mAP指标的字典
    """
    map_scores = {}
    
    for iou_th in iou_thresholds:
        precision, _ = precision_recall(predictions, ground_truths, classes, iou_th)
        ap_scores = {}
        
        for c in classes:
            prec = np.mean(precision[c]) if precision[c] else 0
            ap_scores[c] = prec
            
        map_scores[f'mAP@{iou_th}'] = np.mean(list(ap_scores.values())) if ap_scores else 0
        
    return map_scores


class AverageMeter:
    """
    用于计算和存储运行平均值和当前值。
    """
    
    def __init__(self):
        self.reset()

    def reset(self):
        """重置所有统计信息。"""
        self.val = 0  # 当前值 (last value)
        self.avg = 0  # 平均值 (average)
        self.sum = 0  # 累加和 (sum)
        self.count = 0  # 计数值 (count)

    def update(self, val, n=1):
        """更新统计信息。
        
        Args:
            val: 要添加的新值。
            n: 与此值关联的样本数量（默认为1）。
                例如，如果 `val` 是一个批次的平均损失，`n` 应该是该批次的样本数。
        """
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count